#coding=utf8
# Copyright (c) 2016 Tinydot. inc.
# All Rights Reserved.
#
#    Licensed under the Apache License, Version 2.0 (the "License"); you may
#    not use this file except in compliance with the License. You may obtain
#    a copy of the License at
#
#         http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
#    License for the specific language governing permissions and limitations
#    under the License.


import torch
from .instanceseg import create_stack_pos_map
from mmcv.ops.saconv import SAConv2d
SAConv2d = SAConv2d


class PosConv2d(torch.nn.Module):

    def __init__(self, in_channels, out_channels, kernel_size, **kw):
        super().__init__()
        conv_cls = kw.get("conv_cls", torch.nn.Conv2d)
        self.norm = kw.get("norm_pos", True)
        for k in ["conv_cls", "norm_pos"]:
            if k in kw:
                del kw["conv_cls"]
        self.conv = conv_cls(in_channels+2, out_channels, kernel_size, **kw)
        self.pos_cache = {}

    def create_pos(self, shape, device):
        key = "%s-%s" % (shape, device)
        if key not in self.pos_cache:
            b, _, w, h = shape
            self.pos_cache[key] = torch.from_numpy(create_stack_pos_map(w,
                                                                        h,
                                                                        b, norm=self.norm)
                                                   ).permute(0, 3, 1, 2).to(device)
        return self.pos_cache[key]

    def forward(self, x):
        device = x.device
        return self.conv(torch.cat([x, self.create_pos(x.shape, device).type(x.dtype)], dim=1))

